import os
import shutil
import numpy as np
import onnxruntime as ort
from PIL import Image
from torchvision import transforms


def load_classes(train_dir):
    """加载类别信息"""
    return sorted(os.listdir(train_dir))


def preprocess_image(image_path):
    """预处理图像"""
    transform = transforms.Compose(
        [
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )
    img = Image.open(image_path).convert("RGB")
    return transform(img).unsqueeze(0).numpy()


def predict_with_onnx(model_path, image_paths, classes, threshold=0.5):
    """使用ONNX模型进行预测"""
    # 创建ONNX运行时会话
    session = ort.InferenceSession(model_path)

    results = []
    for img_path in image_paths:
        try:
            # 预处理图像
            input_data = preprocess_image(img_path)

            # 进行预测
            outputs = session.run(None, {"input": input_data})[0]
            probs = np.exp(outputs) / np.sum(np.exp(outputs), axis=1, keepdims=True)
            max_prob = np.max(probs)
            pred_idx = np.argmax(probs)

            # 判断是否为未知类别
            if max_prob < threshold:
                results.append(
                    {
                        "image": img_path,
                        "prediction": "unknown",
                        "probability": float(max_prob),
                    }
                )
            else:
                results.append(
                    {
                        "image": img_path,
                        "prediction": classes[pred_idx],
                        "probability": float(max_prob),
                    }
                )
        except Exception as e:
            results.append({"image": img_path, "prediction": "error", "error": str(e)})

    return results


if __name__ == "__main__":

    output_dir = 'C:\MrDoc\二次跑unknown'
    # 配置
    model_path = "model.onnx"
    train_dir = "data/TrainingSet"

    # 加载类别
    classes = load_classes(train_dir)

    # 获取要预测的图片路径
    input_path = input("请输入要预测的图片路径或目录：")

    # 如果是目录，获取所有图片文件
    if os.path.isdir(input_path):
        image_paths = []
        for root, _, files in os.walk(input_path):
            for file in files:
                if file.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".gif")):
                    image_paths.append(os.path.join(root, file))
    else:
        image_paths = [input_path]

    if not image_paths:
        print("未找到任何图片文件")
        exit()

    # 进行预测
    results = predict_with_onnx(model_path, image_paths, classes)

    # 输出结果
    print("\n预测结果：")
    with open('output.txt', 'w', encoding='utf-8') as f:
        for result in results:
            if "error" in result:
                f.write(f"图片: {result['image']}, 错误: {result['error']}\n")
            else:
                f.write(
                    f"图片: {result['image']}, 预测类别: {result['prediction']}, 置信度: {result['probability']:.4f}\n"
                )
                    # 创建预测类别的文件夹
                class_dir = os.path.join(output_dir, str(result['prediction']))
                os.makedirs(class_dir, exist_ok=True)
                # 创建unknowns文件夹
                unknown_dir = os.path.join(output_dir, 'unknown')
                os.makedirs(unknown_dir, exist_ok=True)
                shutil.copy(result['image'], class_dir)
                # if result['probability'] > 0.9:
                #     # 复制图像到相应的文件夹
                #     shutil.copy(result['image'], class_dir)
                # else:
                #     shutil.copy(result['image'], unknown_dir)